import cv2
import xml.etree.ElementTree as ET
import os


def parse_voc_annotation(annotation_file):
    """解析VOC格式的标注文件，返回包含边界框和标签名的列表"""
    tree = ET.parse(annotation_file)
    root = tree.getroot()

    boxes_with_labels = []
    for obj in root.iter('object'):
        bbox = obj.find('bndbox')
        x1 = int(bbox.find('xmin').text)
        y1 = int(bbox.find('ymin').text)
        x2 = int(bbox.find('xmax').text)
        y2 = int(bbox.find('ymax').text)
        label = obj.find('name').text
        boxes_with_labels.append((x1, y1, x2, y2, label))
    return boxes_with_labels


def visualize_annotations(image_folder, annotation_folder):
    """读取图片和标签，将标签结果显示在图片上"""
    for filename in os.listdir(image_folder):
        if filename.endswith(('.jpg', '.jpeg', '.png')):
            image_path = os.path.join(image_folder, filename)
            # 构造对应的标注文件名
            annotation_filename = os.path.splitext(filename)[0] + '.xml'
            annotation_path = os.path.join(annotation_folder, annotation_filename)

            if os.path.exists(annotation_path):  # 确保标注文件存在
                image = cv2.imread(image_path)
                boxes = parse_voc_annotation(annotation_path)

                for (x1, y1, x2, y2, label) in boxes:
                    # 绿色边框
                    cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    # 在边界框上方显示标签
                    cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0),2)

                cv2.imshow('Annotated Image', image)
                cv2.waitKey(0)
                cv2.destroyAllWindows()


# 图片和标签文件夹路径
imgs_folder = 'imgs'
annotations_folder = 'annotation'

# 可视化显示图片和标签
visualize_annotations(imgs_folder, annotations_folder)